import pygame
import sys
import random

# 初始化 pygame
pygame.init()

# 设置棋盘大小
WIDTH, HEIGHT = 600, 600
GRID_SIZE = 15  # 棋盘 15x15
CELL_SIZE = WIDTH // GRID_SIZE

# 颜色设置
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
GRAY = (169, 169, 169)

# 创建屏幕
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("五子棋 - AI 对战")

# 游戏状态
board = [[0] * GRID_SIZE for _ in range(GRID_SIZE)]  # 0代表空，1代表黑，2代表白
current_player = 1  # 1代表黑棋，2代表白棋


# 检查是否有玩家胜利
def check_win(player):
    for x in range(GRID_SIZE):
        for y in range(GRID_SIZE):
            if board[x][y] == player:
                # 横向
                if y + 4 < GRID_SIZE and all(board[x][y + i] == player for i in range(5)):
                    return True
                # 纵向
                if x + 4 < GRID_SIZE and all(board[x + i][y] == player for i in range(5)):
                    return True
                # 主对角线
                if x + 4 < GRID_SIZE and y + 4 < GRID_SIZE and all(board[x + i][y + i] == player for i in range(5)):
                    return True
                # 副对角线
                if x + 4 < GRID_SIZE and y - 4 >= 0 and all(board[x + i][y - i] == player for i in range(5)):
                    return True
    return False


# 绘制棋盘和棋子
def draw_board():
    screen.fill(WHITE)

    # 绘制网格
    for i in range(GRID_SIZE):
        pygame.draw.line(screen, GRAY, (i * CELL_SIZE, 0), (i * CELL_SIZE, HEIGHT), 1)
        pygame.draw.line(screen, GRAY, (0, i * CELL_SIZE), (WIDTH, i * CELL_SIZE), 1)

    # 绘制已下的棋子
    for x in range(GRID_SIZE):
        for y in range(GRID_SIZE):
            if board[x][y] == 1:  # 黑棋
                pygame.draw.circle(screen, BLACK, (y * CELL_SIZE + CELL_SIZE // 2, x * CELL_SIZE + CELL_SIZE // 2),
                                   CELL_SIZE // 2 - 2)
            elif board[x][y] == 2:  # 白棋
                pygame.draw.circle(screen, WHITE, (y * CELL_SIZE + CELL_SIZE // 2, x * CELL_SIZE + CELL_SIZE // 2),
                                   CELL_SIZE // 2 - 2)

    pygame.display.flip()  # 更新屏幕


# AI 评估函数
def evaluate_board():
    # 简单的评分机制：计算连续相同棋子的数量
    score = 0
    for x in range(GRID_SIZE):
        for y in range(GRID_SIZE):
            if board[x][y] != 0:
                player = board[x][y]
                score += evaluate_position(x, y, player)
    return score


# 评估每个位置的得分
def evaluate_position(x, y, player):
    score = 0
    directions = [((0, 1), (1, 0)), ((1, 0), (1, 1)), ((1, 0), (1, -1))]  # 横、纵、两个对角
    for direction in directions:
        count = 1
        for dx, dy in direction:
            nx, ny = x + dx, y + dy
            if 0 <= nx < GRID_SIZE and 0 <= ny < GRID_SIZE and board[nx][ny] == player:
                count += 1
            else:
                break
        score += count
    return score


# Minimax 算法与 Alpha-Beta 剪枝
def minimax(depth, alpha, beta, maximizing_player):
    if depth == 0 or check_win(1) or check_win(2):
        return evaluate_board()

    if maximizing_player:
        max_eval = -float('inf')
        for x in range(GRID_SIZE):
            for y in range(GRID_SIZE):
                if board[x][y] == 0:
                    board[x][y] = 2  # AI 落子
                    eval = minimax(depth - 1, alpha, beta, False)
                    board[x][y] = 0  # 撤销
                    max_eval = max(max_eval, eval)
                    alpha = max(alpha, eval)
                    if beta <= alpha:
                        break  # 剪枝
        return max_eval
    else:
        min_eval = float('inf')
        for x in range(GRID_SIZE):
            for y in range(GRID_SIZE):
                if board[x][y] == 0:
                    board[x][y] = 1  # 玩家落子
                    eval = minimax(depth - 1, alpha, beta, True)
                    board[x][y] = 0  # 撤销
                    min_eval = min(min_eval, eval)
                    beta = min(beta, eval)
                    if beta <= alpha:
                        break  # 剪枝
        return min_eval


# AI 选择最佳落子
def ai_move():
    best_move = None
    best_value = -float('inf')

    for x in range(GRID_SIZE):
        for y in range(GRID_SIZE):
            if board[x][y] == 0:
                board[x][y] = 2  # AI 落子
                move_value = minimax(2, -float('inf'), float('inf'), False)  # 设置搜索深度为 2
                board[x][y] = 0  # 撤销
                if move_value > best_value:
                    best_value = move_value
                    best_move = (x, y)
    return best_move


# 游戏主循环
def game_loop():
    global current_player

    running = True
    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            elif event.type == pygame.MOUSEBUTTONDOWN:
                # 获取鼠标点击位置
                x, y = event.pos
                row, col = y // CELL_SIZE, x // CELL_SIZE

                # 如果点击的位置是空的，且是玩家的回合
                if board[row][col] == 0 and current_player == 1:  # 只允许黑棋玩家下棋
                    board[row][col] = current_player

                    # 检查是否有玩家胜利
                    if check_win(current_player):
                        print(f"玩家 {current_player} 获胜!")
                        running = False

                    # 切换玩家
                    current_player = 2  # AI 轮到下棋

        # 如果是 AI 的回合
        if current_player == 2:
            ai_x, ai_y = ai_move()
            board[ai_x][ai_y] = current_player
            print(f"AI 下棋: ({ai_x}, {ai_y})")

            # 检查是否有玩家胜利
            if check_win(current_player):
                print(f"玩家 {current_player} 获胜!")
                running = False

            # 切换玩家
            current_player = 1  # 切换回玩家轮到下棋

        # 绘制棋盘
        draw_board()

    pygame.quit()
    sys.exit()


# 启动游戏
if __name__ == "__main__":
    game_loop()
